from abc import ABC, abstractmethod
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier
import pickle
import os

class BaseModel(ABC):
    """
    模型基类
    """
    def __init__(self):
        self.model = None
    
    @abstractmethod
    def train(self, X, y):
        """
        训练模型
        """
        pass
    
    @abstractmethod
    def predict(self, X):
        """
        预测
        """
        pass
    
    def save_model(self, filepath):
        """
        保存模型
        """
        with open(filepath, 'wb') as f:
            pickle.dump(self, f)
    
    @classmethod
    def load_model(cls, filepath):
        """
        加载模型
        """
        with open(filepath, 'rb') as f:
            return pickle.load(f)

class RandomForestModel(BaseModel):
    """
    随机森林模型
    """
    def __init__(self, n_estimators=100, random_state=42):
        super().__init__()
        self.model = RandomForestClassifier(n_estimators=n_estimators, random_state=random_state)
    
    def train(self, X, y):
        self.model.fit(X, y)
    
    def predict(self, X):
        return self.model.predict_proba(X)[:, 1]

class LogisticRegressionModel(BaseModel):
    """
    逻辑回归模型
    """
    def __init__(self, C=1.0, max_iter=1000, random_state=42):
        super().__init__()
        self.model = LogisticRegression(C=C, max_iter=max_iter, random_state=random_state)
    
    def train(self, X, y):
        self.model.fit(X, y)
    
    def predict(self, X):
        return self.model.predict_proba(X)[:, 1]

class SVMModel(BaseModel):
    """
    支持向量机模型
    """
    def __init__(self, probability=True, random_state=42):
        super().__init__()
        self.model = SVC(probability=probability, random_state=random_state)
    
    def train(self, X, y):
        self.model.fit(X, y)
    
    def predict(self, X):
        return self.model.predict_proba(X)[:, 1]

class NeuralNetworkModel(BaseModel):
    """
    神经网络模型
    """
    def __init__(self, hidden_layer_sizes=(100,), max_iter=500, random_state=42):
        super().__init__()
        self.model = MLPClassifier(hidden_layer_sizes=hidden_layer_sizes, max_iter=max_iter, random_state=random_state)
    
    def train(self, X, y):
        self.model.fit(X, y)
    
    def predict(self, X):
        return self.model.predict_proba(X)[:, 1]

class ModelFactory:
    """
    模型工厂类
    """
    @staticmethod
    def create_model(model_type, **kwargs):
        """
        创建模型
        
        Parameters:
        model_type: str, 模型类型 ('rf', 'lr', 'svm', 'nn')
        **kwargs: 模型参数
        
        Returns:
        BaseModel: 模型实例
        """
        if model_type == 'rf':
            return RandomForestModel(**kwargs)
        elif model_type == 'lr':
            return LogisticRegressionModel(**kwargs)
        elif model_type == 'svm':
            return SVMModel(**kwargs)
        elif model_type == 'nn':
            return NeuralNetworkModel(**kwargs)
        else:
            raise ValueError(f"Unsupported model type: {model_type}")